package org.example.partition;

import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.example.flink.pojo.WaterSensor;

import java.util.Arrays;

/**
 * @Author: tang
 * @Description:
 * @Date 2025/2/18 10:03
 */
public class FlinkPartition {

    public static void main(String[] args) throws Exception{

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());

        environment.setParallelism(4);

        DataStreamSource<String> streamSource = environment.socketTextStream("127.0.0.1", 7777);

//        streamSource
//                .rescale()
//                .print();

        environment.fromCollection(Arrays.asList(
                        new WaterSensor("1", "11", 20),
                        new WaterSensor("2", "22", 20),
                        new WaterSensor("2", "2222", 22),
                        new WaterSensor("3", "33", 20)
                ))
                .partitionCustom(new Partitioner<String>() {
                    @Override
                    public int partition(String key, int numPartitions) {
                        int num = key.hashCode() % numPartitions;
                        System.out.println("Partitions---> " + num);
                        return num;
                    }
                }, (KeySelector<WaterSensor, String>) value -> value.getId())
                .print()
                ;

        environment.execute();

    }

}
